"""
@author: 江同学呀
@file: filter_search_results.py
@date: 2025/1/8 23:38
@desc:
    https://www.elastic.co/guide/en/elasticsearch/reference/7.17/filter-search-results.html
"""
from typing import Union, Literal, Dict

from espc.common.common import Number
from espc.common.query_common import ScoreMode
from espc.orm.model.base.base import _Base
from espc.orm.model.dsl.queries.base_queries import _BaseQueries


class Rescorer(_Base):
    """
    查询重新评分器
    https://www.elastic.co/guide/en/elasticsearch/reference/7.17/filter-search-results.html#query-rescorer

    The query rescorer executes a second query only on the Top-K results returned by the query and post_filter phases.
    The number of docs which will be examined on each shard can be controlled by the window_size parameter, which
    defaults to 10.
    By default the scores from the original query and the rescore query are combined linearly to produce the final
    _score for each document. The relative importance of the original query and of the rescore query can be controlled
    with the query_weight and rescore_query_weight respectively. Both default to 1.

    查询重新评分器仅对查询和 post_filter 阶段返回的 Top-K 结果执行第二个查询。将在每个分片上检查的文档数量可由 window_size 参数控制，默认为 10。
    默认情况下，原始查询和 rescore 查询的分数以线性方式组合，以生成每个文档的最终_score。原始查询和 rescore 查询的相对重要性可以分别使用
    query_weight 和 rescore_query_weight 来控制。两者都默认为 1。

    :param window_size:
    :param score_mode:
    :param rescore_query:
    :param query_weight:
    :param rescore_query_weight:
    """

    def __init__(
            self, window_size: int, score_mode: Union[Literal["multiply", "total", "avg", "max", "min"], ScoreMode],
            rescore_query: _BaseQueries = None, query_weight: Number = None, rescore_query_weight: Number = None
    ):
        super().__init__()
        self._window_size: int = window_size
        self._score_mode: Union[Literal["multiply", "total", "avg", "max", "min"], ScoreMode, None] = score_mode
        self._rescore_query: _BaseQueries = rescore_query
        self._query_weight: Number = query_weight
        self._rescore_query_weight: Number = rescore_query_weight
        return

    def _build(self) -> Dict:
        _body: Dict = {
            "rescore_query":
                self._rescore_query
                if isinstance(self._rescore_query, Dict)
                else {
                    self._rescore_query.type: self._rescore_query._build()
                },
        }
        if self._score_mode:
            _body["score_mode"] = self._score_mode if isinstance(self._score_mode, str) else self._score_mode.value
        if self._query_weight:
            _body["query_weight"] = self._query_weight
        if self._rescore_query_weight:
            _body["rescore_query_weight"] = self._rescore_query_weight
        return {
            "window_size": self._window_size,
            "query": _body
        }
